MODULE fefmm_density

   USE fmm_global_paras
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: fefmm_init_density,   &
             fefmm_get_density,    &
             fefmm_get_projAO

   TYPE(fmm_basis),    POINTER, SAVE :: basis
   TYPE(fmm_sh_pairs), POINTER, SAVE :: sh_pairs(:)
   REAL(REALK),        POINTER, SAVE :: Dmat(:,:)

   REAL(REALK), ALLOCATABLE, TARGET, SAVE :: projAO(:,:)

   INTEGER(INTK) :: FEdim
   REAL(REALK)   :: FEgrain
   REAL(REALK)   :: box_grain

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_init_density(scheme,basis_tar,sh_pairs_tar,Dmat_tar)

      IMPLICIT NONE
      TYPE(scheme_paras),         INTENT(IN) :: scheme
      TYPE(fmm_basis),    TARGET, INTENT(IN) :: basis_tar
      TYPE(fmm_sh_pairs), TARGET, INTENT(IN) :: sh_pairs_tar(:)
      REAL(REALK),        TARGET, INTENT(IN) :: Dmat_tar(:,:)

      NULLIFY(basis, sh_pairs, Dmat)
      basis => basis_tar
      sh_pairs => sh_pairs_tar(:)
      Dmat => Dmat_tar(:,:)

      FEdim   = scheme%FEdim   ! number of grid points in one box (1D)

      FEgrain = scheme%grain / (FEdim-1)
      box_grain = scheme%grain
      
      ALLOCATE (projAO(FEdim**3,basis%nbas))

   END SUBROUTINE fefmm_init_density

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_free_density

      CONTINUE

   END SUBROUTINE fefmm_free_density

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_get_density(box_cntr,dens)

      USE fefmm_box_utils, ONLY: fefmm_sh_pair_in_box

      IMPLICIT NONE
      REAL(REALK), INTENT(IN)  :: box_cntr(3)
      REAL(REALK), INTENT(OUT) :: dens(:,:,:)

      REAL(REALK)   :: tmp
      INTEGER(INTK) :: i,j,k,ij,ijk, p,q, Ish,Jsh
      INTEGER(INTK) :: IAnglA,IAnglB, IL1,IL2,IL2Temp
      LOGICAL       :: IEqJ

      CALL fefmm_project_AOs(basis,box_cntr)
!      DO ij = 1, SIZE(sh_pairs)
!         IF (.NOT. fefmm_sh_pair_in_box(sh_pairs(ij),box_cntr,box_grain) ) &
!            CYCLE
!         CALL fefmm_project_AOs(basis,sh_pairs(ij)%J,box_cntr)
!      END DO

      dens = zero

      DO ij = 1, SIZE(sh_pairs)

         IF (.NOT. fefmm_sh_pair_in_box(sh_pairs(ij),box_cntr) ) &
            CYCLE

         Ish = sh_pairs(ij)%I
         Jsh = sh_pairs(ij)%J
         IAnglA = basis%KType(Ish)
         IAnglB = basis%KType(Jsh) 
         IEqJ = (ISh == JSh)

         p = basis%KLoc_Car(ISh)
         DO IL1 = basis%LtuvMin_Car(IAnglA), basis%LtuvMax_Car(IAnglA)
            p = p+1
            q = basis%KLoc_Car(JSh)

            IL2Temp = basis%LtuvMax_Car(IAnglB)
            IF (IEqJ) IL2Temp = IL1
            DO IL2 = basis%LtuvMin_Car(IAnglB), IL2Temp
               q = q+1
               IF (ABS(Dmat(p,q)) < 1d-15) CYCLE

               ijk = 0
               DO i = 1, FEdim
                  DO j = 1, FEdim
                     DO k = 1, FEdim
                        ijk = ijk +1
                        tmp =  Dmat(p,q) * projAO(ijk,p) * projAO(ijk,q)
                        dens(k,j,i) = dens(k,j,i) + tmp
                     END DO 
                  END DO 
               END DO 

            END DO 
         END DO 

      END DO 

   END SUBROUTINE fefmm_get_density

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_project_AOs(basis,box_cntr)

      IMPLICIT NONE
      TYPE(fmm_basis), INTENT(IN) :: basis
      REAL(REALK),     INTENT(IN) :: box_cntr(3)

      REAL(REALK)   :: origin(3), gpt(3), AOcntr(3)
      REAL(REALK)   :: ExpF, X,Y,Z, R2, Xn,Yn,Zn
      INTEGER(INTK) :: ijk, p, t,u,v
      INTEGER(INTK) :: Ish, IAnglA, IL1 
      INTEGER(INTK) :: npx, npy, npz

      origin(:) = box_cntr(:) - half*box_grain

      shells: DO Ish = 1, basis%nshells
         IAnglA = basis%KType(Ish)
         AOcntr(:) = basis%Centr(:,basis%KAtom(Ish))
   
         p = basis%KLoc_Car(ISh)
         AOs: DO IL1 = basis%LtuvMin_Car(IAnglA), basis%LtuvMax_Car(IAnglA)
            p = p+1
            projAO(:,p) = zero
            t = basis%Lt(IL1)
            u = basis%Lu(IL1)
            v = basis%Lv(IL1)
   
            ijk = 0
            DO npz = 1, FEdim
               gpt(3) = origin(3) + (npz-1)*FEgrain
               Z  = (gpt(3) - AOcntr(3))
               Zn = Z**v
   
               DO npy = 1, FEdim
                  gpt(2) = origin(2) + (npy-1)*FEgrain
                  Y  = (gpt(2) - AOcntr(2))
                  Yn = Y**u
   
                  DO npx = 1, FEdim
                     ijk = ijk +1
                     gpt(1) = origin(1) + (npx-1)*FEgrain
                     X  = (gpt(1) - AOcntr(1))
                     Xn = X**t
   
                     R2 = X*X + Y*Y + Z*Z
                     CALL fefmm_get_ExpF(basis,Ish,R2,ExpF)
   
                     projAO(ijk,p) = projAO(ijk,p) + ExpF*Xn*Yn*Zn
   
                  END DO
               END DO
   
            END DO
         END DO AOs
      END DO shells


   END SUBROUTINE fefmm_project_AOs

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_get_ExpF(basis,Ish,R2,ExpF)

      IMPLICIT NONE
      TYPE(fmm_basis), INTENT(IN)  :: basis
      INTEGER(INTK),   INTENT(IN)  :: Ish
      REAL(REALK),     INTENT(IN)  :: R2
      REAL(REALK),     INTENT(OUT) :: ExpF

      INTEGER(INTK) :: i, ip1, ip2
      REAL(REALK)   :: ExpA

      ExpF = zero

      ip1 = basis%KStart(Ish)
      ip2 = ip1 + basis%KontG(Ish) - 1
      DO i = ip1, ip2
         ExpA = basis%Expnt(i)
         ExpF = ExpF + basis%CCoef(i)*EXP(-ExpA*R2)
      END DO

   END SUBROUTINE fefmm_get_ExpF

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_get_projAO(p,projAO_ptr)

      IMPLICIT NONE
      INTEGER(INTK),        INTENT(IN)  :: p
      REAL(REALK), POINTER, INTENT(OUT) :: projAO_ptr(:)

      projAO_ptr => projAO(:,p)

   END SUBROUTINE fefmm_get_projAO

!-------------------------------------------------------------------------------

END MODULE fefmm_density
